"""
Multi-armed Bandits问题 - ε-贪心算法实现
文件: epsilon_greedy.py
描述: 实现书中第2章介绍的ε-贪心算法，用于解决Multi-armed Bandits问题

作者: [您的名字]
日期: 2023-11-10
"""

import numpy as np
import matplotlib.pyplot as plt
from typing import List, Tuple, Dict, Optional, Union
import time
import argparse
import os

class MultiArmedBandit:
    """
    Multi-armed Bandits环境类
    
    属性:
        k (int): tiger机的臂数
        true_values (np.ndarray): 每个臂的真实期望奖励值
        best_action (int): 具有最高真实期望奖励的臂的索引
        variance (float): 奖励分布的方差
    """
    
    def __init__(self, k: int = 10, mean: float = 0.0, variance: float = 1.0):
        """
        初始化Multi-armed Bandits环境
        
        参数:
            k (int): tiger机的臂数，默认为10
            mean (float): 臂的真实奖励均值的初始均值，默认为0
            variance (float): 臂的真实奖励均值的方差，默认为1
        """
        self.k = k
        # 从正态分布采样每个臂的真实期望奖励
        self.true_values = np.random.normal(mean, np.sqrt(variance), k)
        # 找出具有最高真实期望奖励的臂
        self.best_action = np.argmax(self.true_values)
        # 奖励分布的方差
        self.variance = variance
        
    def pull(self, action: int) -> float:
        """
        拉动指定的臂，返回奖励
        
        参数:
            action (int): 要拉动的臂的索引（从0到k-1）
            
        返回:
            float: 从N(true_value[action], 1)分布中采样的奖励值
        """
        # 检查动作有效性
        if action < 0 or action >= self.k:
            raise ValueError(f"动作必须在0到{self.k-1}之间，收到{action}")
            
        # 返回带有噪声的奖励（真实期望值加上标准正态噪声）
        return np.random.normal(self.true_values[action], 1.0)
    
    def reset(self) -> None:
        """
        重置tiger机环境，重新采样每个臂的真实期望奖励
        """
        self.true_values = np.random.normal(0, 1, self.k)
        self.best_action = np.argmax(self.true_values)
        
    def get_optimal_action(self) -> int:
        """
        返回具有最高真实期望奖励的臂的索引
        
        返回:
            int: 最优臂的索引
        """
        return self.best_action


class EpsilonGreedy:
    """
    ε-贪心算法实现
    
    属性:
        bandit (MultiArmedBandit): Multi-armed Bandits环境
        epsilon (float): 探索概率
        q_values (np.ndarray): 每个臂的估计价值
        action_counts (np.ndarray): 每个臂被选择的次数
        t (int): 当前时间步
    """
    
    def __init__(self, bandit: MultiArmedBandit, epsilon: float = 0.1, 
                 initial_value: float = 0.0, step_size: Optional[float] = None):
        """
        初始化ε-贪心智能体
        
        参数:
            bandit (MultiArmedBandit): Multi-armed Bandits环境
            epsilon (float): 探索概率，默认为0.1
            initial_value (float): 价值估计的初始值，默认为0
            step_size (Optional[float]): 固定步长，如果为None则使用1/n
        """
        self.bandit = bandit
        self.epsilon = epsilon
        self.k = bandit.k
        self.step_size = step_size
        
        # 初始化每个臂的估计价值
        self.q_values = np.ones(self.k) * initial_value
        # 记录每个臂被选择的次数
        self.action_counts = np.zeros(self.k, dtype=int)
        # 当前时间步
        self.t = 0
        
        # 保存历史数据用于分析
        self.reward_history = []
        self.action_history = []
        self.optimal_action_history = []
        
    def select_action(self) -> int:
        """
        根据ε-贪心策略选择动作
        
        返回:
            int: 选择的臂的索引
        """
        # 以epsilon的概率随机选择一个臂（探索）
        if np.random.random() < self.epsilon:
            return np.random.randint(self.k)
        
        # 以1-epsilon的概率选择估计价值最高的臂（利用）
        # 处理多个臂具有相同最大估计价值的情况
        max_value = np.max(self.q_values)
        max_actions = np.where(self.q_values == max_value)[0]
        return np.random.choice(max_actions)
    
    def update(self, action: int, reward: float) -> None:
        """
        根据获得的奖励更新动作价值估计
        
        参数:
            action (int): 选择的臂的索引
            reward (float): 获得的奖励
        """
        # 更新时间步
        self.t += 1
        
        # 更新该臂被选择的次数
        self.action_counts[action] += 1
        
        # 确定步长
        alpha = self.step_size if self.step_size is not None else 1.0 / self.action_counts[action]
        
        # 使用增量更新公式更新估计价值
        self.q_values[action] += alpha * (reward - self.q_values[action])
        
        # 记录历史数据
        self.reward_history.append(reward)
        self.action_history.append(action)
        self.optimal_action_history.append(1 if action == self.bandit.get_optimal_action() else 0)
        
    def run(self, n_steps: int = 1000) -> Tuple[List[float], List[int], List[int]]:
        """
        运行智能体n_steps步
        
        参数:
            n_steps (int): 要运行的步数，默认为1000
            
        返回:
            Tuple[List[float], List[int], List[int]]: 
                (奖励历史, 动作历史, 最优动作选择历史)
        """
        # 重置历史数据
        self.reward_history = []
        self.action_history = []
        self.optimal_action_history = []
        
        for _ in range(n_steps):
            # 选择动作
            action = self.select_action()
            
            # 执行动作获取奖励
            reward = self.bandit.pull(action)
            
            # 更新估计
            self.update(action, reward)
            
        return self.reward_history, self.action_history, self.optimal_action_history


def run_experiment(n_runs: int = 2000, n_steps: int = 1000, k: int = 10, 
                   epsilons: List[float] = [0, 0.01, 0.1, 0.2]) -> Dict[str, Dict[str, np.ndarray]]:
    """
    运行多次实验，比较不同ε值的性能
    
    参数:
        n_runs (int): 每个ε值运行的次数
        n_steps (int): 每次运行的步数
        k (int): tiger机的臂数
        epsilons (List[float]): 要测试的ε值列表
        
    返回:
        Dict[str, Dict[str, np.ndarray]]: 
            包含平均奖励和最优动作选择比例的结果字典
    """
    results = {
        'average_rewards': {},
        'optimal_action_pct': {}
    }
    
    for epsilon in epsilons:
        print(f"运行实验 ε = {epsilon}")
        
        rewards = np.zeros((n_runs, n_steps))
        optimal_actions = np.zeros((n_runs, n_steps))
        
        start_time = time.time()
        
        for i in range(n_runs):
            # 每次运行使用新的tiger机
            bandit = MultiArmedBandit(k=k)
            agent = EpsilonGreedy(bandit, epsilon=epsilon)
            
            # 运行智能体
            run_rewards, _, run_optimal_actions = agent.run(n_steps)
            
            # 记录结果
            rewards[i] = run_rewards
            optimal_actions[i] = run_optimal_actions
            
            # 显示进度
            if (i + 1) % 100 == 0:
                elapsed = time.time() - start_time
                remaining = (elapsed / (i + 1)) * (n_runs - i - 1)
                print(f"  完成 {i+1}/{n_runs} 次运行 "
                      f"(用时: {elapsed:.1f}s, 预计剩余: {remaining:.1f}s)")
        
        # 计算平均结果
        results['average_rewards'][f"ε={epsilon}"] = np.mean(rewards, axis=0)
        results['optimal_action_pct'][f"ε={epsilon}"] = np.mean(optimal_actions, axis=0) * 100
        
    return results


def plot_results(results: Dict[str, Dict[str, np.ndarray]], save_dir: str = None) -> None:
    """
    绘制实验结果图表
    
    参数:
        results (Dict[str, Dict[str, np.ndarray]]): 实验结果
        save_dir (str): 保存图表的目录，如果为None则显示图表
    """
    plt.figure(figsize=(15, 6))
    
    # 绘制平均奖励
    plt.subplot(1, 2, 1)
    for label, values in results['average_rewards'].items():
        plt.plot(values, label=label)
    plt.xlabel('步数')
    plt.ylabel('平均奖励')
    plt.title('不同ε值的平均奖励')
    plt.legend()
    plt.grid(alpha=0.3)
    
    # 绘制最优动作百分比
    plt.subplot(1, 2, 2)
    for label, values in results['optimal_action_pct'].items():
        plt.plot(values, label=label)
    plt.xlabel('步数')
    plt.ylabel('选择最优动作的百分比 (%)')
    plt.title('不同ε值选择最优动作的百分比')
    plt.legend()
    plt.grid(alpha=0.3)
    
    plt.tight_layout()
    
    # 保存或显示图表
    if save_dir:
        os.makedirs(save_dir, exist_ok=True)
        plt.savefig(os.path.join(save_dir, 'epsilon_greedy_comparison.png'), dpi=300)
        print(f"图表已保存到 {os.path.join(save_dir, 'epsilon_greedy_comparison.png')}")
    else:
        plt.show()


if __name__ == "__main__":
    # 解析命令行参数
    parser = argparse.ArgumentParser(description='Multi-armed Bandits问题 - ε-贪心算法实验')
    parser.add_argument('--runs', type=int, default=2000, help='运行次数')
    parser.add_argument('--steps', type=int, default=1000, help='每次运行的步数')
    parser.add_argument('--k', type=int, default=10, help='tiger机的臂数')
    parser.add_argument('--epsilons', type=float, nargs='+', default=[0, 0.01, 0.1, 0.2],
                        help='要测试的ε值列表')
    parser.add_argument('--save-dir', type=str, default=None, 
                        help='保存图表的目录，如果不指定则显示图表')
    args = parser.parse_args()
    
    print("Multi-armed Bandits问题 - ε-贪心算法实验")
    print("=" * 50)
    print(f"运行次数: {args.runs}")
    print(f"每次步数: {args.steps}")
    print(f"tiger机臂数: {args.k}")
    print(f"探索率(ε)值: {args.epsilons}")
    print(f"图表保存目录: {args.save_dir if args.save_dir else '不保存，直接显示'}")
    print("=" * 50)
    
    # 运行实验
    results = run_experiment(
        n_runs=args.runs, 
        n_steps=args.steps, 
        k=args.k, 
        epsilons=args.epsilons
    )
    
    # 绘制并保存/显示结果
    plot_results(results, save_dir=args.save_dir)
    
    print("实验完成！")